import pandas as pd
from sklearn.metrics import f1_score, accuracy_score
# from simpletransformers.classification import ClassificationModel
from simpletransformers.model import TransformerModel

# 注意如果用windows运行程序的话，要加if __name__ == '__main__':才能正常运行
if __name__ == '__main__':
# Train and Evaluation data needs to be in a Pandas Dataframe of two columns. The first column is the text with type str, and the second column is the label with type int.
    train_data = [
        ["物流速度快，物美价廉", 1],
        ["快又好用", 1],
        ["就是快递太慢了", 0],
        ["盒子没贴封条，袋子没有封条，袋子都磨损了，京东怎么部倒闭呢", 0],
        ["发货速度挺快的",1],
        ["一直京东购物,服务好,物流快",1],
        ["物流快，血压计有背光和语音，老人家用非常方便",1],
        ["预约15点前送到，最后17:30才送到",0],
    ]
    train_df = pd.DataFrame(train_data)
    # print(train_df)

    eval_data = [
        ["物流速度不行", 0],
        ["物流画了四五天才寄到，太慢了", 0],
        ["东西还没送到，电话不打一下私自放在快递柜里，定单显示本人已经签收了，京东服务越来越差了",0],
        ["快递及时送到，很不错啊",1],
    ]
    eval_df = pd.DataFrame(eval_data)


# Create a ClassificationModel
#   选择的预训练模型"hfl/chinese-roberta-wwm-ext"，程序会自动到https://huggingface.co/上去下载模型，不过要和roberta-base模型格式一样的模型,程序下载好了之后才能识别并运行
#   和roberta-base模型格式一样的模型:就是包含的文件数量和命名差不多的才能运行
#   roberta-base模型的格式：https://huggingface.co/roberta-base/tree/main
#   同理hfl/chinese-roberta-wwm-ext模型的格式：https://huggingface.co/hfl/chinese-roberta-wwm-ext/tree/main
#     model = ClassificationModel("roberta", "hfl/chinese-roberta-wwm-ext", use_cuda=False)
    model = TransformerModel('bigbird', 'google/bigbird-pegasus-large-arxiv', use_cuda=False, num_labels=4, args={'learning_rate':1e-5,
                                                                                             'num_train_epochs': 2,
                                                                                             'reprocess_input_data': True,
                                                                                             'overwrite_output_dir': True})
# Train the model
    model.train_model(train_df)

# Evaluate the model
    def f1_multiclass(labels, preds):
        return f1_score(labels, preds, average='micro')
    # result, model_outputs, wrong_predictions = model.eval_model(eval_df)
    result, model_outputs, wrong_predictions = model.eval_model(eval_df, f1=f1_multiclass, acc=accuracy_score)

    predictions, raw_outputs = model.predict(['第十天就到，太慢的了'])

    print(result)
    print(model_outputs)
    print(wrong_predictions)
    print(predictions)
    print(raw_outputs)